package com.zhaiker.object_detection_tnn

import android.content.Context
import android.graphics.Bitmap
import android.util.Log
import android.widget.Toast
import com.cv.tnn.model.Detector
import com.cv.tnn.model.FrameInfo
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import java.io.File
import java.io.FileOutputStream
import java.io.IOException

object GenderAgeHelper {

    private const val DET_MODEL: String = "face/rfb_landm_face_416_416.sim"

    private const val CLS_MODEL: String = "age_gender/age_gender_112_112.sim"

    private var TNN_MODEL_FILES: Array<String> = arrayOf(
        DET_MODEL,
        CLS_MODEL,
    )
    private const val MODEL_ID = 0
    private const val NUM_THREAD = 1
    private const val USE_GPU = false

    private var _sdkAvailable : Boolean = false

    //sdk是否可用
    val sdkAvailable get() = _sdkAvailable

    private suspend fun copyModelFromAssetsToData(context: Context) : Boolean{
        // assets目录下的模型文件名
        //Toast.makeText(App.instance, "Copy model to data...", Toast.LENGTH_SHORT).show()
        try {
            for (tnn_model in TNN_MODEL_FILES) {
                //  {"face/rfb1.0_face_320_320.opt.tnnproto", "face/rfb1.0_face_320_320.opt.tnnmodel"};
                val tnnproto = "$tnn_model.tnnproto"
                val tnnmodel = "$tnn_model.tnnmodel"
                //Log.w(com.cv.tnn.activity.WelcomeActivity.TAG, "copy model:$tnn_model")
                copyAssetFileToFiles(context, tnnproto)
                copyAssetFileToFiles(context, tnnmodel)
            }
            Log.i("object-detection-tnn","Copy tnn model Success")
            //Toast.makeText(context, "Copy model Success", Toast.LENGTH_SHORT).show()
            return true
        } catch (e: Exception) {
            e.printStackTrace()
            Log.i("object-detection-tnn","Copy tnn model Error")
            Toast.makeText(context, "Copy tnn model Error", Toast.LENGTH_SHORT).show()
            return false
        }
    }

    @Throws(IOException::class)
    private suspend fun copyAssetFileToFiles(context: Context, filename: String) = withContext(Dispatchers.IO){
        val file = File(context.filesDir.toString() + File.separator + filename)
        if(file.exists()){
            Log.i("object-detection-tnn","${filename} exists, no copy required")
            return@withContext
        }

        val parent = File(file.parent)
        if (!parent.exists()) {
            parent.mkdirs()
        }
        file.createNewFile()

        val inputStream = context.assets.open(filename)
        val buffer = ByteArray(inputStream.available())
        inputStream.read(buffer)
        inputStream.close()


        val os = FileOutputStream(file)
        os.write(buffer)
        os.close()
    }

    private fun initModel(context : Context) {
        val root: String = context.getFilesDir().toString() + File.separator
        //String tnnproto = WelcomeActivity.TNN_MODEL_FILES[MODEL_ID] + ".tnnproto";
        //String tnnmodel = WelcomeActivity.TNN_MODEL_FILES[MODEL_ID] + ".tnnmodel";
        val length: Int = TNN_MODEL_FILES.size
        val det_model: String = TNN_MODEL_FILES[MODEL_ID] // 检测模型
        val cls_model: String = TNN_MODEL_FILES[length - 1] //默认最后一个是识别模型
        Detector.init(
            det_model,
            cls_model,
            root,
            MODEL_ID,
            NUM_THREAD,
            USE_GPU
        )
    }

    suspend fun init(context: Context){
        if(!copyModelFromAssetsToData(context)){
            return
        }
        initModel(context)
        _sdkAvailable = true
    }

    /**
     * 检查性别年龄和人脸box
     * @param bitmap
     * @param score_thresh 置信度阈值
     * @param iou_thresh IOU阈值
     */
    fun detect(bitmap : Bitmap, score_thresh : Float = 0.5f, iou_thresh : Float = 0.5f) : Array<FrameInfo>{
        return Detector.detect(bitmap,score_thresh,iou_thresh)
    }

}